package ArithmeticObjects;

import java.util.Arrays;

import utils.ByteArrUtils;
import utils.ByteTree;

import crypto.HashFunction;

/**
 * Representing a modular group, a subgroup Gq of the multiplicative group Zp*
 * @author Idan
 */
public class ModularGroup extends Group{

	private LargeNumberObject p; //order of the multiplicative group
	private LargeNumberObject q; //order of the sub group
	private ModularGroupElement g; //standard generator
	private LargeNumber pVal;
	private LargeNumber qVal;
	private LargeNumber gVal;	
	private int e; //indicator for encode/decode
	
	
	/**
	 * creates a ModularGroup represented as a ByteTree
	 * @param bt ByteTree representing the ModularGroup
	 * @throws Exception in case the Group's construction failed
	 */
	public ModularGroup(ByteTree bt) throws Exception{
		super(new LargeNumberObject(bt.getChild(0)));
		this.p = new LargeNumberObject(bt.getChild(0));
		this.q = new LargeNumberObject(bt.getChild(1));
		this.g = new ModularGroupElement(bt.getChild(2),this);
		this.e = bt.getChild(3).getNumber().intValue();
		
		this.pVal = p.getValue();
		this.qVal = q.getValue();
		this.gVal = g.getValue();
				
		VerifyGroup();
	}
	
	/**
	 * a method that verified that the Modular Group is valid according to the specification
	 * @throws Exception in case the ModularGroup is incompatible
	 */
	private void VerifyGroup() throws Exception
	{
		//verifying that the orders are prime
		int certainty = (int)Math.pow(10,6);
		if (!p.getValue().isProbablePrime(certainty))
			throw new Exception("P is not prime!");
		if (!q.getValue().isProbablePrime(certainty))
			throw new Exception("Q is not prime!");
		//verifying that p is greater than or equal to 3
		if (pVal.compareTo(new LargeNumber(3)) < 0)
			throw new Exception("Modular group is incompatable: P must be larger or equal to 3");
		if (e == 1)
		{
			//in case e==1, verify that p=2q+1
			if (pVal.compareTo(qVal.multiply(LargeNumber.TWO).add(LargeNumber.ONE)) != 0)
				throw new Exception("Modular group is incompatable: e equals 1, but p<>2q+1");
		}
		else if (e == 2)
		{
			/*in case e==2, verify that p=tq+1 for some integer t and the difference
			between the bit lengths of p and q is less than 2^10*/
			if ((pVal.subtract(LargeNumber.ONE).remainder(qVal).intValue() != 0) ||
					(Math.abs(pVal.bitLength() - qVal.bitLength()) >= Math.pow(2,10)))
				throw new Exception("Modular group is incompatable: e equals 2, but either p<>tq+1 " +
						"or the difference between the bit lengths of p and q is less than 2^10");
		}
		else if (e != 0) throw new Exception("Invalid e value: " + e); //verify that e is in {0, 1, 2}
	}
	
	/**
	 * a helper method to compute the ByteTree message returned by the Decode method
	 * @param n number to decode
	 * @param b bytes to decode
	 * @return ByteTree of the decoded message
	 * @throws Exception in case the decoding failed
	 */
	public ByteTree helper(LargeNumber n, int b) throws Exception
	{
		//calculate k' as 2^(8(b+4))
		LargeNumber kPrime = LargeNumber.TWO.power(new LargeNumber(8).multiply(new LargeNumber(b+4)));
		//calculating temp as n mod k' which is interpreted as bytes_4(l) | m'
		LargeNumber temp = n.mod(kPrime);
		byte[] lTemp = new byte[4];
		byte[] mPrime = new byte[b];
		byte[] tTemp = temp.toByteArray();
		byte[] tTemp1 = new byte[b+4];
		System.arraycopy(tTemp, 0, tTemp1, b+4-tTemp.length, tTemp.length);
		System.arraycopy(tTemp1, 0, lTemp, 0, 4); //bytes_4(l)
		System.arraycopy(tTemp1, 4, mPrime, 0, b); //m'
		int l = new LargeNumber(lTemp).intValue();
		//if l<0 or l>b we return an empty message
		if (l<0 || l>b) 
		{
			l=0;
			return null;
		}
		byte[] m = new byte[l];
		System.arraycopy(mPrime, 0, m, 0, l);
		//return m as the first l bytes of m'
		return ByteTree.CreateLeaf(m);
	}
	
	/**
	 * encodes a message represented by a ByteTree to a ModularGroupElement
	 * @param m the message represented as a ByteTree
	 * @return a ModularGroupElement representing the encoded message
	 * @throws Exception in case the encoding failed
	 */
	public ModularGroupElement Encode(ByteTree m) throws Exception{
		int mSize = m.getSize(); //bm as an int
		HashFunction h = new HashFunction("SHA-256");
		LargeNumber bm = new LargeNumber(mSize);
		LargeNumber t = pVal.subtract(LargeNumber.ONE).divide(qVal);
		
		int n = pVal.bitLength();
		int b = (int)Math.floor(((double)n - 2) / 8) - 4;
		
		
		if (e == 0 ) {
			byte[] temp = new byte[mSize + 1];
			LargeNumber l;
			LargeNumber a = gVal;
			//Going over all of the group members of Q to find an element that represents 
			//the encoded message m. 
			for (LargeNumber i = LargeNumber.ZERO; i.compareTo(qVal.subtract(LargeNumber.ONE)) < 0; 
					i.add(LargeNumber.ONE))
			{
				//extracting the first bm+1 bytes of H(a)
				System.arraycopy(h.digest(a.toByteArray()), 0, temp, 0, mSize + 1);
				l = new LargeNumber(temp);
				System.arraycopy(temp, 1 , temp, 0, mSize + 1);
				
				//checking if the first bm+1 bytes of H(a) are of the form
				//bytes_1(l) | m for some l such that l mod 4 = bm.
				if ((Arrays.equals(temp, m.getRawByteTree())) &&
						l.mod(new LargeNumber(4)).compareTo(bm) == 0)
					return new ModularGroupElement(ByteTree.CreateLeaf((a.toByteArray())), this);
				else a = a.multiply(gVal);
			}
			return null;
		}
		else if (e == 1) {
			
			byte[] mPrime = new byte[b];
			byte[] k = new byte[mSize+4];
			//If bm = 0, we set m' to 01 | bytes_b-1(0)
			if (mSize == 0)
			{
				mPrime[0] = 1;
				for (int i=1; i < b; i++)
					mPrime[i] = 0;
			}
			//otherwise we set m' to m | bytes_b-bm(0)
			else {
				System.arraycopy(m, 0, mPrime, 0, mSize);
				for (int i = mSize+1; i < b; i++)
					mPrime[i] = 0;
			}
			System.arraycopy(ByteArrUtils.intToByteArray(mSize), 0, k, 0, 4);
			System.arraycopy(mPrime, 0, k, 5, b);
			LargeNumber a = gVal;
			LargeNumber kVal = new LargeNumber(k);
			
			//Going over all of the group members of Q to find an element that represents 
			//the encoded message m. 
			for (LargeNumber i = LargeNumber.ONE; i.compareTo(qVal.subtract(LargeNumber.ONE)) < 0; 
					i.add(LargeNumber.ONE))
			{
				//if a = k or p - k return it
				if (a.compareTo(kVal) == 0 || a.compareTo(pVal.subtract(kVal)) == 0)
					return new ModularGroupElement(ByteTree.CreateLeaf((a.toByteArray())), this);
				else a = a.multiply(gVal);
			}
			return null;
		}
		else if (e == 2) {
			LargeNumber nPrime;
			
			//calculating n' = ceiling(n/8)+1
			if (t.remainder(new LargeNumber(8)).compareTo(new LargeNumber(4)) > 0)
				nPrime = (t.divide(new LargeNumber(8)).add(LargeNumber.TWO));
			else nPrime = (t.divide(new LargeNumber(8)).add(LargeNumber.ONE));
			
			b = (int)Math.floor(n / 8) - nPrime.intValue() - 4;
			
			//interpret k as the integer representation of bytes_4(bm) | m | bytes_b-bm(0)
			byte[] kArr = new byte[b+4];
			System.arraycopy(ByteArrUtils.intToByteArray(mSize), 0, kArr, 0, 4);
			System.arraycopy(m, 0, kArr, 5, mSize);
			for (int i=mSize+1; i<b+4; i++)
				kArr[i] = 0;
			LargeNumber k = new LargeNumber(kArr);
			LargeNumber a = gVal;
			LargeNumber kPrime = LargeNumber.TWO.power(new LargeNumber(8).multiply(new LargeNumber(b+4)));
			LargeNumber i = LargeNumber.ZERO;
			
			//Going over all of the group members of Q to find an element that represents 
			//the encoded message m. 
			for (LargeNumber j = LargeNumber.ZERO; j.compareTo(qVal.subtract(LargeNumber.ONE)) < 0; 
					j.add(LargeNumber.ONE))
			{
				//if a = i*2^(8(b+4))+k return it
				if (a.compareTo(i.multiply(kPrime).add(k)) == 0)
					return new ModularGroupElement(ByteTree.CreateLeaf((a.toByteArray())), this);
				else {
					a = a.multiply(gVal);
					i.add(LargeNumber.ONE);
				}
			}
		}
		return null;	
	}
	
	/**
	 * decodes a ModularGroupElement to a message represented by a ByteTree 
	 * @param a a ModularGroupElement representing the message to decode
	 * @return a ByteTree representing the decoded message
	 * @throws Exception in case the decoding failed
	 */
	public ByteTree Decode(ModularGroupElement a) throws Exception {
		HashFunction h = new HashFunction("SHA-256");
		LargeNumber aVal = a.getValue();
		int n = pVal.bitLength();
		int b = (int)Math.floor(((double)n - 2) / 8) - 4;
		if (e == 0)
		{
			byte[] d = h.digest(aVal.toByteArray());
			//interpreting l' as the first byte of H(a)
			LargeNumber lPrime = new LargeNumber(d[0]);
			//l = l' mod 4
			int l = lPrime.mod(new LargeNumber(4)).intValue();
			byte[] m = new byte[l];
			System.arraycopy(d, 1, m, 0, l);
			//returning the next l bytes of H(a)
			return ByteTree.CreateLeaf(m);
		}
		else if (e == 1)
		{
			LargeNumber k;
			//defining k as a or p-a depending on if a<p-a or not
			if (a.getValue().compareTo(pVal.subtract(aVal)) < 0)
				k = aVal;
			else k = pVal.subtract(aVal);
			//using a helper method to return the first l bytes of m'
			//where k mod 2^(8(b+4)) = bytes_4(l) | m'
			return helper(k, b);
		}
		else if (e == 2)
		{
			//using a helper method to return the first l bytes of m'
			//where a mod 2^(8(b+4)) = bytes_4(l) | m'
			return helper(aVal, b);
		}
		return null;
	}
	
	/**
	 * @return the sub group's order
	 */
	public Element getQ() {
		return q;
	}
	
	/**
	 * @return the multiplicative group's order
	 */
	public Element getP() {
		return p;
	}
	
	/**
	 * @return the standard generator
	 */
	public Element getG() {
		return g;
	}
	
	/**
	 * @return the indicator e
	 */
	public int getE() {
		return e;
	}
	
	/**
	 * compares the current instance to another ModularGroup
	 * @param group the ModularGroup to compare to
	 * @return 0 if they are the same and 1 otherwise
	 * @throws Exception in case the comparison failed
	 */
	public int compareTo(ModularGroup group) throws Exception{
		if (p.compareTo(group.getP()) != 0) return 1;
		if (q.compareTo(group.getQ()) != 0) return 1;
		if (g.compareTo(group.getG()) != 0) return 1;
		if (e != group.getE()) return 1;
		return 0;
	}
}
